import os
import random
from functools import partial
from multiprocessing.pool import Pool
from pathlib import Path
from random import sample

import cv2
import pandas as pd
from tqdm import tqdm

from constants import DFDC
from option import parse_args

os.environ["MKL_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["OMP_NUM_THREADS"] = "1"

cv2.ocl.setUseOpenCL(False)
cv2.setNumThreads(0)


def cmp_frame_name(frame_name):
    frame_name_split = frame_name.split('_')
    return int(frame_name_split[0])


def get_real_fake_pairs(data_root, mode, sample_num=0):
    video_names = open(os.path.join(data_root, '{}.txt'.format(mode)), 'r').readlines()
    real_fake_pairs = []
    for video_name in video_names:
        video_name_split = video_name.strip().split(',')
        pair = (video_name_split[0], video_name_split[1], video_name_split[2])
        real_fake_pairs.append(pair)
    if sample_num:
        real_fake_pairs = sample(real_fake_pairs, sample_num)
    return real_fake_pairs


def get_paths(vid, label, root_dir, sample_num: int):
    part_id, ori_id, fake_id = vid
    ori_dir = os.path.join(root_dir, 'crops', part_id, ori_id)
    fake_dir = os.path.join(root_dir, 'crops', part_id, fake_id)
    mask_dir = os.path.join(root_dir, 'diffs', part_id, fake_id)

    data = []
    # Some unaligned videos have been removed
    if label == 1 and not os.path.exists(fake_dir):
        return data
    # Some masks may have been removed
    if label == 1 and not os.path.exists(mask_dir):
        return data

    frames = os.listdir(fake_dir)
    frames.sort(key=cmp_frame_name)
    for frame in frames:
        frame_split = str(frame).split('_')
        frame_num = int(frame_split[0])
        if frame_num == 0 or frame_num % 10 != 0:
            continue
        if frame_num // 10 > sample_num:
            break
        ori_img_path = os.path.join(ori_dir, frame)
        fake_img_path = os.path.join(fake_dir, frame)
        if label == 0 and os.path.exists(ori_img_path):
            data.append([ori_img_path, label, ori_id])
        elif label == 1 and os.path.exists(fake_img_path) and os.path.exists(ori_img_path):
            mask_path = os.path.join(mask_dir, "{}_diff.png".format(frame[:-4]))
            if os.path.exists(mask_path):
                data.append([fake_img_path, label, ori_id])

    return data


def save_folds(args, save_path, pairs_mode, mode, sample_num=100):
    data = []
    if mode == 'train_real':
        ori_ori_pairs = set([(part, ori, ori) for part, ori, fake in pairs_mode])
        ori_fake_pairs = set()
    elif mode == 'train_fake':
        ori_ori_pairs = set()
        ori_fake_pairs = set(pairs_mode)
    else:
        ori_ori_pairs = set([(part, ori, ori) for part, ori, fake in pairs_mode])
        ori_fake_pairs = set(pairs_mode)

    ori_ori_pairs = list(ori_ori_pairs)
    ori_fake_pairs = list(ori_fake_pairs)
    ori_ori_pairs.sort(key=cmp)
    ori_fake_pairs.sort(key=cmp)

    with Pool(processes=1) as p:
        # original label=0
        with tqdm(total=len(ori_ori_pairs)) as pbar:
            func = partial(get_paths, label=0, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_ori_pairs):
                pbar.update()
                data.extend(v)
        # fake label=1
        with tqdm(total=len(ori_fake_pairs)) as pbar:
            func = partial(get_paths, label=1, root_dir=args.root_dir, sample_num=sample_num)
            for v in p.imap_unordered(func, ori_fake_pairs):
                pbar.update()
                data.extend(v)
    fold_data = []
    for img_path, label, ori_vid in data:
        path = Path(img_path)
        part = path.parent.parent.name
        video = path.parent.name
        file = path.name
        fold_data.append([os.path.join(part, video), file, label, os.path.join(part, ori_vid), int(file.split("_")[0])])
    # random.shuffle(fold_data)
    columns = ["video", "file", "label", "original", "frame"]
    file_name = 'data_{}_{}.csv'.format(DFDC, mode)
    save_file = os.path.join(save_path, file_name)
    pd.DataFrame(fold_data, columns=columns).to_csv(save_file, index=False)


def real_fake_split(ori_fake_pairs):
    ori_fake_pairs_real = set()
    ori_fake_pairs_fake = set()
    for pair in ori_fake_pairs:
        ori_fake_pairs_fake.add(pair)
        ori_fake_pairs_real.add((pair[0], pair[1], pair[1]))
    return ori_fake_pairs_real, ori_fake_pairs_fake


def cmp(pair):
    pair_part_split = pair[0].split('_')
    return '{:02d}'.format(int(pair_part_split[3])) + pair[1] + pair[2]


def main():
    random.seed(110)
    args = parse_args()
    save_path = '../data/{}'.format(DFDC)
    os.makedirs(save_path, exist_ok=True)

    # original_path = os.path.join(args.root_dir, 'original_sequences', 'youtube', 'c23')
    # manipulated_path = os.path.join(args.root_dir, 'manipulated_sequences', args.fake_type, 'c23')
    pairs_train = get_real_fake_pairs(args.root_dir, 'train')
    pairs_val = get_real_fake_pairs(args.root_dir, 'val')
    pairs_test = get_real_fake_pairs(args.root_dir, 'test')
    print(len(pairs_train), len(pairs_val), len(pairs_test))

    sample_num_train = 10
    sample_num_val = 10
    sample_num_test = 20
    check_mask = True

    save_folds(args, save_path, pairs_train, 'train', sample_num=sample_num_train)
    save_folds(args, save_path, pairs_val, 'val', sample_num=sample_num_val)
    save_folds(args, save_path, pairs_test, 'test', sample_num=sample_num_test)

    ori_fake_pairs_real, ori_fake_pairs_fake = real_fake_split(pairs_train)
    save_folds(args, save_path, ori_fake_pairs_real, 'train_real', sample_num=sample_num_train)
    save_folds(args, save_path, ori_fake_pairs_fake, 'train_fake', sample_num=sample_num_train)


if __name__ == '__main__':
    # Total: 2000
    main()
